from pylab import *
from TDS import TDS
from dispCal.disp import calDisp
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 定义模型参数
z = np.array([0, 0.1, 0.2, 0.5, 2])
thickness = z[1:]-z[:-1]
vs = np.array([0.4, 0.6, 0.8, 1.0])      # 横波速度 (km/s)
vp = np.array([1.5, 2.0, 2.3, 2.5])      # 纵波速度 (km/s)
qs = np.array([500.0, 500.0, 500.0, 500.0])  # 横波品质因子
qp = np.array([1000.0, 1000.0, 1000.0, 1000.0])  # 纵波品质因子
rho = np.array([2.0, 2.0, 2.0, 2.0])  # pn密度
no = 100
nv = 100
print(vp.min(),vp.max())
# 创建TDS实例并设置参数
startT = time.time()

tds = TDS()

# 设置媒质参数
tds.set_media_parameters(
    z=z,
    vp0=vp,
    vs0=vs,
    rho=rho,
    qs=qs,
    qp=qp,
)

# 设置源参数和计算参数
tds.set_source_parameters(
    zs=0.08, z0=0.0, 
    mino=0.1, maxo=7.0, no=no,
    minv=vs.min()/2, maxv=np.max(vp), nv=nv)

tds.set_computation_parameters(m=11, Twin=30000.0)

tds.write_all_parameters_in_fortran()
tds.save_parameters("grt.dat")


# 构建频率和速度数组用于显示
freqs = np.linspace(tds.source['mino'], tds.source['maxo'], tds.source['no'])
vels = np.linspace(tds.source['minv'], tds.source['maxv'], tds.source['nv'])


F,V = np.meshgrid(freqs, vels)
F = F.flatten()
V = V.flatten()
# 计算频散谱（底图）
print("开始计算频散谱底图...")
# spectrum = tds.calculate_dispersion_spectrum()
spectrum = tds.calculate_dispersion_spectrum_arrays(F,V).reshape([nv,no])
endT = time.time()
print(f"频散谱计算完成，耗时: {endT-startT:.2f} 秒")
spectrum = np.abs(spectrum)
# for i in range(spectrum.shape[0]):
#     spectrum[i] = spectrum[i] / spectrum[i].max()
print(f"频散谱计算完成，结果形状: {spectrum.shape}")

spec1 = np.copy(spectrum)
for j in range(no):
    sj = spec1[j,:]
    sj = sj/np.max(sj)
    sj = -np.diff(sj, 2)
    # sj = np.sign(np.abs())
    # sj = np.sign(sj-np.max(sj)/10)+1

    spec1[j,1:nv-1] = sj
spec1[:,0]=0
spec1[:,-1]=0


# 使用dispCal计算理论频散曲线
print("计算理论频散曲线...")

# 绘制频散谱底图和理论曲线
fig = figure(figsize=(6,4), dpi=600)

# 绘制频散谱底图
subplot(121)
pcolormesh(freqs,vels, spectrum, cmap='nipy_spectral_r',shading='nearest')
colorbar(label='频散谱幅度')
clim([0, spectrum.max()/10])
# clim([0, 100])
# 叠加理论频散曲线
for order in range(1, 3):
    velocities = calDisp(thickness, vp, vs, rho, 1/freqs[-1::-1], wave='rayleigh', mode=order, 
                         velocity='phase', flat_earth=True, ar=6370, dc0=0.005,
                         Qs=qs, Qp=qp)[-1::-1]
    plot(freqs,velocities, 'k-', linewidth=0.3, label=f'阶数 {order}')

ylabel('相速度 (km/s)')
xlabel('频率 (Hz)')
title('频散谱底图与理论频散曲线对比')
legend()
ylim(tds.source['minv'], tds.source['maxv'])
xlim(tds.source['mino'], tds.source['maxo'])

subplot(122)
pcolormesh(freqs,vels, spec1, cmap='nipy_spectral_r',shading='nearest')


fig.tight_layout()
fig.savefig('dispersion_comparison.png', bbox_inches='tight')

print("结果已保存到 dispersion_comparison.png")